import os
import sys
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '../..'))
sys.path.append(project_root)
import pickle
import argparse
import matplotlib
matplotlib.use('Agg')

import numpy as np

from PIL import Image
from utils import *

X_RANGE = None
Y_RANGE = None

def load_label(label_path):
    label = Image.open(label_path)
    label = np.array(label, dtype=np.float32) / 255.0
    return label

def plot_point_and_mask(frame_data, label, output_path):
    X = np.array(frame_data['x'])
    Y = np.array(frame_data['y'])
    rgb = np.array(frame_data['rgb'])

    resolution = 900
    x_min, x_max = X_RANGE
    y_min, y_max = Y_RANGE

    x_scale = (x_max - x_min) / resolution
    y_scale = (y_max - y_min) / resolution

    rgb_image = np.zeros((resolution, resolution, 3), dtype=np.float32)
    count_matrix = np.zeros((resolution, resolution), dtype=np.int32)

    for i in range(len(X)):
        pixel_x = int((X[i] - x_min) / x_scale)
        pixel_y = int((y_max - Y[i]) / y_scale)
        if 0 <= pixel_x < resolution and 0 <= pixel_y < resolution:
            rgb_image[pixel_y, pixel_x] += rgb[i]
            count_matrix[pixel_y, pixel_x] += 1

    non_zero_mask = count_matrix > 0
    rgb_image[non_zero_mask] /= count_matrix[non_zero_mask][:, None]

    label_resized = np.array(Image.fromarray((label * 255).astype(np.uint8)).resize((resolution, resolution), Image.NEAREST)) / 255.0
    label_overlay = np.stack([label_resized] * 3, axis=-1)
    combined_image = rgb_image * 0.8 + label_overlay * 0.2

    output_image = (combined_image * 255).astype(np.uint8)
    Image.fromarray(output_image).save(output_path)

def process_pkl_file(pkl_file_path, label_folder, output_folder):
    with open(pkl_file_path, 'rb') as f:
        data_dict = pickle.load(f)

    all_frames = data_dict.get('data', [])
    X_list, Y_list, Z_list = [], [], []
    DN_list, azi_list, ele_list, err_list = [], [], [], []

    for frame in all_frames:
        if len(frame) == 0:
            continue
        X_list.append(frame['x'])
        Y_list.append(frame['y'])
        Z_list.append(frame['z'])
        DN_list.append(frame['dn'])
        azi_list.append(frame['aziLinearlityMap'])
        ele_list.append(frame['eleLinearlityMap'])
        err_list.append(frame['errorMap'])

    X = np.concatenate(X_list) if X_list else np.array([])
    Y = np.concatenate(Y_list) if Y_list else np.array([])
    Z = np.concatenate(Z_list) if Z_list else np.array([])
    DN = np.concatenate(DN_list) if DN_list else np.array([])
    azi = np.concatenate(azi_list) if azi_list else np.array([])
    ele = np.concatenate(ele_list) if ele_list else np.array([])
    err = np.concatenate(err_list) if err_list else np.array([])

    raw_data = {
        'x': X,
        'y': Y,
        'z': Z,
        'dn': DN,
        'aziLinearlityMap': azi,
        'eleLinearlityMap': ele,
        'errorMap': err
    }

    filtered_data = filter_z(raw_data)
    filtered_data = filter_data(filtered_data, X_RANGE, Y_RANGE, filter_noise=True)

    X = filtered_data['x']
    Y = filtered_data['y']
    Z = filtered_data['z']

    DN = filtered_data['dn']
    azi = filtered_data['aziLinearlityMap']
    ele = filtered_data['eleLinearlityMap']
    err = filtered_data['errorMap']
    
    s = azi * ele * err

    rgb = compute_rgb(Z, s, DN)

    frame_data = {
        'x': X,
        'y': Y,
        'rgb': rgb
    }

    base_name = os.path.splitext(os.path.basename(pkl_file_path))[0]
    label_path = os.path.join(label_folder, f"{base_name}.png")
    if not os.path.exists(label_path):
        print(f"找不到对应的 label 文件: {label_path}, 将跳过叠加。")
        label = np.zeros((512, 512))
    else:
        label = load_label(label_path)

    output_path = os.path.join(output_folder, f"{base_name}.png")
    plot_point_and_mask(frame_data, label, output_path)

def main():
    global X_RANGE, Y_RANGE

    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument("--X_RANGE", type=str, required=True)
    parser.add_argument("--Y_RANGE", type=str, required=True)
    args = parser.parse_args()
    
    X_RANGE = eval(args.X_RANGE)
    Y_RANGE = eval(args.Y_RANGE) 
    
    pkl_folder = "./workspace/image/SAR/input"
    label_folder = "./workspace/image/SAR/label"
    show_dir = "./workspace/image/SAR/show"

    os.makedirs(show_dir, exist_ok=True)
    
    all_files = []
    with os.scandir(pkl_folder) as entries:
        files = [entry for entry in entries if entry.is_file() and entry.name.endswith(".pkl")]
        files.sort(key=lambda entry: natural_key(entry.name))
        
        for file in files:
            all_files.append(file.path)

    total_files = len(all_files)
    for index, file_path in enumerate(all_files, start=1):
        print(f"正在处理文件 ({index}/{total_files})")
        process_pkl_file(file_path, label_folder, show_dir)

if __name__ == "__main__":
    main()
